import torch
import torch.nn as nn
from AI.nets.AttU_Net import AttU_net
import torchvision.transforms as transforms
from torchvision.utils import *
from PIL import Image
import cv2
import numpy as np
class AI_Segment():
    def __init__(self):

        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        self.model = AttU_net(1,1)
        self.check_path = './AI/weights/best_seg.pth'
        # map_location = lambda storage, loc: storage
        self.model.load_state_dict(torch.load(self.check_path, map_location=self.device))
        self.model.to(self.device)
        self.model.eval()

    def tensor2numpy(self,input_tensor: torch.Tensor):
        input_tensor = input_tensor.to(torch.device('cpu')).detach().numpy()
        # in_arr = np.transpose(input_tensor, (
        # 1, 2, 0))  # 将(c,w,h)转换为(w,h,c)。但此时如果是全精度模型，转化出来的dtype=float64 范围是[0,1]。后续要转换为cv2对象，需要乘以255
        cvimg =np.uint8(input_tensor * 255)
        return cvimg


    def start_seg(self,image):
        w,h = image.shape[0:2]
        # img_merge = cv2.merge([image,image,image])
        image_copy = image.copy()
        img_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        img_pil = img_pil.resize((128, 128), Image.BICUBIC)
        transform_list = transforms.Compose([transforms.Grayscale(1),
                                             transforms.Resize(128),
                                             transforms.ToTensor()])
        img = transform_list(img_pil)
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        img = img.to(self.device)
        outputs = self.model(img)
        out_img = outputs.view(outputs.shape[0], 1, 128, 128)
        print("output")
        save_image(out_img, 'test.jpg')
        grid = make_grid(out_img)
        ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
        img_ori = cv2.resize(ndarr,(h,w),cv2.INTER_CUBIC)
        _, img_1 = cv2.threshold(img_ori[:,:,0], 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # cv2.imshow('b', img_1)
        contours,hierarchy = cv2.findContours(img_1,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        img_contour = cv2.drawContours(image_copy,contours,-1,(0,0,255),4)
        # cv2.imshow('a',img_contour)
        # cv2.waitKey(0)
        im = Image.fromarray(ndarr)
        return img_contour